"""
强化学习智能体算法模块

该模块实现了多智能体强化学习算法的主控制器，负责智能体的决策制定、奖励生成、
模型训练和保存/加载等功能。继承自RLAlgorithmBase基类，实现了完整的强化学习流程。

主要功能：
1. 智能体决策制定和动作选择
2. 自定义奖励函数设计
3. 模型训练和更新
4. 模型保存和加载
5. 轨迹管理和数据收集

注意：此为红方(Red Team)使用的算法实现，与蓝方版本类似但奖励设计不同。
红方需要到达目标点完成任务。
"""

import os  # 操作系统接口
import torch  # PyTorch深度学习框架
import shutil  # 文件操作工具
from config import GlobalConfig  # 全局配置
from Utils.tensor_ops import repeat_at, gather_righthand  # 张量操作工具
from Common.AlgorithmBase import RLAlgorithmBase  # 强化学习算法基类
import numpy as np  # 数值计算库
from collections import Counter  # 计数器工具
from Common.action_set_encoder import encode_action_as_digits  # 动作编码工具


class AlgorithmConfig:
    """
    强化学习算法配置类

    存储模型训练相关的超参数和配置信息。
    某些配置项会被启动文件中的GlobalConfig或ScenarioConfig覆写。
    """

    num_heads = 8  # 多头注意力机制的头数量(用于注意力网络)
    gamma = 0.99  # 折扣因子，用于计算未来奖励的现值
    tau = 0.95  # GAE(Generalized Advantage Estimation)平滑参数
    # tau值控制优势估计的偏差-方差权衡，较小值偏向即时奖励，较大值考虑更多未来奖励

    num_training_trajectories = 3  # 开始训练前需要采集的轨迹数量
    # 每轮训练对应一条完整的训练轨迹

    use_batch_norm = True  # 是否使用批归一化
    load_trained_model = False  # 是否加载预训练模型
    specific_model_path = ''  # 预训练模型的具体路径

    train_epoch = 16  # 每次训练的轮次
    num_batches = 1  # 将训练数据集划分的批次数量
    max_grad_norm = 0.5  # 梯度裁剪的最大范数，防止梯度爆炸
    learning_rate = 1e-4  # 学习率
    memory_safety_check = False  # 内存/显存安全检查标志
    # 设置为true时，会将数据集划分为多个批次以防止内存溢出

    reward_forwarding = False  # 是否启用奖励转发机制
    # 奖励转发：将"死亡"智能体的未来奖励转发给存活智能体
    reward_forwarding_gamma = 0.99  # 奖励转发的折扣因子

    hidden_dim = 256  # 神经网络隐藏层维度
    num_agent = 'auto load, do not change'  # 智能体数量(由基类自动设置)
    action_filter_enabled = True  # 是否启用可用动作过滤
    # 使用可用动作掩码来限制模型生成的动作，确保只选择有效的动作
    skip_test = True  # 是否跳过测试阶段


class ReinforceAgentAlgorithm(RLAlgorithmBase):
    """
    强化学习智能体算法主类

    继承自RLAlgorithmBase，实现了完整的多智能体强化学习算法流程。
    负责智能体的决策制定、奖励计算、模型训练和模型管理等功能。

    注意：此为红方使用的算法实现，红方的任务目标是到达指定目标点。
    """

    def __init__(self, n_agent, n_thread, space, team=None):
        """
        初始化强化学习算法

        Args:
            n_agent (int): 智能体数量
            n_thread (int): 并行线程数量(环境数量)
            space (object): 空间配置对象
            team (int, optional): 队伍编号
        """
        super().__init__(n_agent, n_thread, space, team)

        # 初始化核心组件
        self.trainer = None  # 训练器对象
        self.policy = None   # 策略网络
        self.agent_list = None  # 智能体列表
        self.n_thread = n_thread  # 线程数量
        self.avail_act = None  # 可用动作列表
        self.opp_agent_not_alive = None  # 对方阵亡智能体信息
        self.opp_agent_uid = None  # 对方智能体UID

        self.previous_distance = None  # 上一时刻的距离(用于距离相关奖励计算)

        # ----------------------RL模型设置区域-------------------------
        # 注意：以下代码已注释，实际的模型初始化在set_algorithm_trainer方法中
        # from .qmix_net import QMIX
        # self.policy = QMIX(num_agents=n_agent,
        #                    state_dim=self.rawob_dim * GlobalConfig.ScenarioConfig.obs_n_entity,
        #                    action_space=self.n_actions, device=self.device)
        # self.policy = self.policy.to(self.device)
        # # ----------------------结束设置RL模型-------------------------
        #
        # # 初始化优化器和轨迹管理器
        # from .train import Trainer
        # from Common.trajectory import TrajectoryManager
        # self.trainer = Trainer(self.policy, train_config=AlgorithmConfig)
        # self.trajectory_pool = TrajectoryManager(
        #     num_env=n_thread, max_trajectory_data=int(GlobalConfig.ScenarioConfig.MaxStepsPerEpisode),
        #     train_model_func=self.trainer.train_model_on_trajectories,team=self.team)
        #
        # if AlgorithmConfig.load_trained_model:
        #     self.load_model()

        # 获取父文件夹名称作为模型名称(2025/07添加)
        current_file_path = os.path.abspath(__file__)
        current_dir_path = os.path.dirname(current_file_path)
        parent_dir_path = os.path.dirname(current_dir_path)
        parent_folder_name = os.path.basename(parent_dir_path)
        self.model_name = parent_folder_name  # 模型名称
        self.logdir = ('./TeamAlg/%s/TrainLogs/' % self.model_name)  # 日志目录

    def set_algorithm_trainer(self):
        """
        设置算法训练器和策略网络

        该方法延迟初始化强化学习相关的组件，包括：
        1. QMIX策略网络
        2. 训练器
        3. 轨迹管理器
        4. 模型加载(如果需要)
        """
        from .qmix_net import QMIX  # 导入QMIX网络

        # 初始化QMIX策略网络
        self.policy = QMIX(num_agents=self.n_agent,
                          state_dim=self.rawob_dim * GlobalConfig.ScenarioConfig.obs_n_entity,
                          action_space=self.n_actions, device=self.device)
        self.policy = self.policy.to(self.device)

        # 初始化训练器和轨迹管理器
        from .train import Trainer  # 导入训练器
        from Common.trajectory import TrajectoryManager  # 导入轨迹管理器
        self.trainer = Trainer(self.policy, train_config=AlgorithmConfig)
        self.trajectory_pool = TrajectoryManager(
            num_env=self.n_thread,
            max_trajectory_data=int(GlobalConfig.ScenarioConfig.MaxStepsPerEpisode),
            train_model_func=self.trainer.train_model_on_trajectories, team=self.team)

        # 如果配置要求加载预训练模型
        if AlgorithmConfig.load_trained_model:
            self.load_model()

    def making_decision(self, state, test_mode):
        """
        智能体决策制定主函数

        根据当前状态选择动作，实现智能体的决策逻辑。

        Args:
            state (dict): 环境状态信息，包含观测数据等
            test_mode (bool): 是否为测试模式

        Returns:
            tuple: (选择的动作数组, 更新后的状态字典)
        """
        # 确保钩子函数已清除
        assert ('_hook_' not in state)

        # 如果是训练模式且不是测试，检查是否需要训练
        if not state['Test-Flag'] and GlobalConfig.train_mode:
            self.train()

        # 从环境状态中提取关键信息
        observation, threads_active_flag, avail_act = \
            tuple(state[item] if item in state else None for item
                  in ('obs', 'threads_active_flag', 'avail_act'))

        # 数据有效性检查
        assert observation is not None and len(observation) == sum(threads_active_flag)
        # 确保可用动作信息正确
        if AlgorithmConfig.action_filter_enabled:
            assert avail_act is not None

        # 使用策略网络选择动作
        action = self.policy.select_action(states=observation,
                                           test_mode=(test_mode and not AlgorithmConfig.skip_test),
                                           avail_act=avail_act)

        # 构建轨迹片段，用于后续训练
        traj_framefrag = {
            "_SKIP_": ~threads_active_flag,  # 跳过非活跃线程
            "avail_act": avail_act,          # 可用动作
            "state": observation,            # 状态
            "action": action,                # 动作
        }
        if avail_act is not None:
            traj_framefrag.update({'avail_act': avail_act})

        # 如果不是测试模式，设置钩子函数用于后续奖励处理
        if not test_mode:
            state['_hook_'] = self.trace_update_callback(traj_framefrag, req_hook=True)

        return action.copy(), state

    def generate_reward(self, state, env):
        """
        自定义奖励函数生成器

        根据游戏事件和智能体状态生成奖励信号。红方的奖励设计主要基于：
        1. 到达目标点的距离变化奖励
        2. 所有智能体都到达目标点的团队奖励
        3. 摧毁敌方单位的奖励
        4. 超出边界区域的负奖励

        Args:
            state (dict): 环境状态信息
            env (object): 环境对象

        Returns:
            numpy.ndarray: 每个智能体的奖励值
        """
        print("Call team-%d algorithm reward function" % self.team)

        # 获取智能体配置信息
        agent_property = GlobalConfig.ScenarioConfig.SubTaskConfig.agent_list
        team_counts = Counter(item['team'] for item in agent_property)

        # 更新当前队伍的智能体列表
        if self.agent_list is None or env.current_step == 1:
            self.agent_list = np.array([agent for agent in env.agents if agent.team == self.team])

        # 初始化奖励数组
        if self.Enable_Reward_Sharing:  # 启用奖励共享：整个队伍共享一个奖励
            reward = np.zeros(shape=env.n_teams)
        else:  # 每个智能体独立的奖励
            reward = np.zeros(shape=team_counts[self.team])

        # 红方(team=0)的奖励设计：基于到达目标点的任务
        if self.team == 0:
            # 获取目标点位置
            target_x = state['dataGlobal']['keyObjArr'][0]['location']['x']
            target_y = state['dataGlobal']['keyObjArr'][0]['location']['y']
            target_position = np.array([target_x, target_y])  # 目标位置

            # 获取红方所有智能体的当前位置
            agent_position = [item.pos2d for item in self.agent_list]
            current_distance_team = [np.linalg.norm(position - target_position) for position in agent_position]

            # 初始化上一时刻距离
            if env.current_step == 1:
                self.previous_distance = current_distance_team

            # 获取存活状态
            f_alive = np.array([agent.alive for agent in self.agent_list])

            # 初始化到达目标标志
            reach_target = np.zeros(shape=(len(self.agent_list)))

            # 为每个智能体计算奖励
            for index, agent in enumerate(self.agent_list):
                if agent.alive:
                    # 超出右边界的负奖励
                    if agent.location[0] > 42830:
                        if not self.Enable_Reward_Sharing:
                            reward[index] = reward[index] - 0.01
                        else:
                            reward[self.team] = reward[self.team] - 0.01

                    # 距离目标点更近的正奖励
                    current_distance = current_distance_team[index]
                    if current_distance < self.previous_distance[index]:
                        self.previous_distance[index] = current_distance
                        if not self.Enable_Reward_Sharing:
                            reward[index] = reward[index] + 0.01
                        else:
                            reward[1] = reward[1] + 0.01  # 注意：这里应该是reward[0]？

                    # 检查是否到达目标点附近(1500范围内)
                    if current_distance < 1500:
                        reach_target[index] = 1

            # 团队奖励：所有存活智能体都到达目标点时给予额外奖励
            all_reached = reach_target[f_alive]
            reached_index = np.where(all_reached)
            if not self.Enable_Reward_Sharing:
                if all(all_reached):
                    reward[reached_index] = reward[reached_index] + 0.15
            else:
                reward[self.team] = reward[self.team] + 0.15  # 红方团队奖励

        # 处理游戏事件，计算摧毁奖励
        events = state['dataGlobal']['events']  # 获取服务端传回的事件列表
        for event in events:
            event_parsed = env.parse_event(event)  # 解析事件信息

            # 处理单位摧毁事件
            if event_parsed['Event'] == 'Destroyed':
                damage_causer_uid = event_parsed['DamageCauser']  # 攻击方UID
                target_id = event_parsed['Target']  # 被摧毁目标UID

                # 根据UID查找对应的智能体实例
                agent_causer = env.find_agent_by_uid(damage_causer_uid)
                agent_target = env.find_agent_by_uid(target_id)

                # 如果攻击方属于蓝方(team=1)，给红方负奖励
                if agent_causer.team == 1:
                    tid = agent_causer.team_id
                    if not self.Enable_Reward_Sharing:
                        reward[tid] += 0.05
                    else:
                        reward[self.team] += 0.05

                # 注释掉的代码：对被摧毁方给予负奖励
                # if agent_target.team == 1:
                #     tid = agent_target.team_id
                #     reward[tid] -= 0.01

        # 返回奖励数组(注释掉的展平操作)
        # reward = reward.flatten(order='C')
        return reward

    def train(self):
        """
        模型训练函数

        检查轨迹池中是否有足够的数据进行训练，如果有则启动训练过程。
        定期保存模型。

        训练触发条件：
        1. 轨迹池中有足够数量的训练轨迹
        2. 每两次训练保存一次模型
        """
        # 检查轨迹池是否准备好进行训练
        if self.trajectory_pool.is_ready_to_train():
            # 使用轨迹数据进行训练
            update_cnt = self.trajectory_pool.train_with_trajectory_data()
            # 每两次训练保存一次模型
            if update_cnt % 2 == 0:
                self.save_model(update_cnt)

    def save_model(self, update_count, info=None):
        """
        保存模型参数

        保存策略网络和优化器的状态字典，用于后续恢复训练。

        触发保存的条件：
        1. 更新次数达到特定值(如50, 100, ...)
        2. 收到外部保存命令(通过info参数)
        3. 检测到保存标志文件

        Args:
            update_count (int): 当前更新次数
            info (str, optional): 附加信息，用于文件命名
        """
        # 创建模型备份目录
        if not os.path.exists('%s/model_backups/' % self.logdir):
            os.makedirs('%s/model_backups/' % self.logdir)

        # 主模型文件路径
        pt_path = '%s/%s_model.pt' % (self.logdir, self.model_name)
        print('Saving %s model to %s' % (self.model_name, pt_path))

        # 准备保存的状态字典：包含策略网络和所有优化器
        save_state = {'policy': self.policy.state_dict()}
        for name, optimizer in self.policy.optimizers_dict.items():
            save_state[f'{name}_state_dict'] = optimizer.state_dict()

        # 保存到主文件
        torch.save(save_state, pt_path)

        # 处理附加信息，生成备份文件名
        if info is None:
            info = str(update_count)
        elif ':' in info:  # 替换冒号以避免文件名问题
            new_info = info.replace(':','-')
            info = ''.join([str(update_count), '_', new_info])
        else:
            info = ''.join([str(update_count), '_', info])

        # 创建备份文件
        pt_path2 = '%s/model_backups/%s_model_%s.pt' % (self.logdir, self.model_name, info)
        shutil.copyfile(pt_path, pt_path2)

        print('[Save Model]: Saving %s model finished!' % self.model_name)

    def load_model(self):
        """
        加载预训练模型

        从指定路径加载模型参数和优化器状态，用于恢复训练或进行推理。
        支持从默认路径或手动指定的路径加载。
        """
        # 确定模型加载路径
        manual_dir = AlgorithmConfig.specific_model_path
        if manual_dir == '':
            ckpt_dir = '%s/%s_model.pt' % (self.logdir, self.model_name)
        else:
            ckpt_dir = '%s/%s' % (self.logdir, manual_dir)

        # 确定设备类型
        cuda_n = 'cpu' if 'cpu' in self.device else self.device
        strict = True  # 严格加载模式

        # 加载检查点文件
        cpt = torch.load(ckpt_dir, map_location=cuda_n)

        # 加载策略网络参数
        self.policy.load_state_dict(cpt['policy'], strict=strict)

        # 加载所有优化器的状态(2024/08添加)
        for name, optimizer in self.policy.optimizers_dict.items():
            optimizer.load_state_dict(cpt[f'{name}_state_dict'])

        print('[Load Model]:', ckpt_dir)